import os
from fastapi import FastAPI
from starlette.websockets import WebSocket, WebSocketDisconnect
import logging

from mcp_server.mcp_client import MCPClient

# WebSocket处理-处理WebSocket连接和消息路由-管理用户会话和上下文
# 配置日志
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# 获取当前文件所在目录的绝对路径
current_dir = os.path.dirname(os.path.abspath(__file__))

async def init_ai_websocket(app: FastAPI):
    # 存储用户上下文的字典
    user_contexts = {}

    @app.websocket("/ws/chat")
    async def websocket_endpoint(websocket: WebSocket):
        user_id = id(websocket) # 为每个用户创建一个唯一的 user_id
        client = None
        try:
            # 1.首先调用 websocket.accept() 接受连接
            await websocket.accept()
            logger.info(f"WebSocket connection accepted for user {user_id}")
            
            # 初始化该用户的上下文
            user_contexts[user_id] = [{"role": "system", "content": "你是一个有帮助的助手。"}]
            
            # 创建并初始化MCP客户端
            client = MCPClient()
            server_path = os.path.join(current_dir, 'mcp_server.py')
            await client.connect_to_server(server_path)
            
            while True:
                try:
                     # 接收WebSocket消息
                    user_msg = await websocket.receive_text()
                    logger.info(f"Received message from user {user_id}: {user_msg}")
                    
                    # 将消息发送回前端显示
                    await websocket.send_json({"role": "user", "content": user_msg})
                    
                    # 将消息发送给AI模型处理-这里调用MCPClient的put_query方法
                    # WebSocket接收消息》调用MCPClient的put_query》put_query调用process_query》process_query处理消息并返回结果》结果通过WebSocket返回给前端
                    response = client.put_query(user_msg)
                    await websocket.send_json({"start": True}) # 发送 start 信号表示开始处理
                    
                    # 通过 async for 循环流式接收AI的回复-将每个回复片段发送给前端
                    async for content_piece in response:
                        await websocket.send_json({"role": "assistant", "content": content_piece})
                    
                    # 发送 done 信号表示处理完成
                    await websocket.send_json({"done": True})
                    
                except Exception as e:
                    logger.error(f"Error processing message for user {user_id}: {str(e)}")
                    await websocket.send_json({
                        "role": "assistant",
                        "content": "抱歉，处理您的消息时遇到问题，请稍后重试。"
                    })
                    
        except WebSocketDisconnect:
            logger.info(f"WebSocket disconnected for user {user_id}")
        except Exception as e:
            logger.error(f"Unexpected error for user {user_id}: {str(e)}")
        finally:
            # 清理资源-在 finally 块中清理用户上下文-关闭MCP客户端连接-记录清理日志
            if user_id in user_contexts:
                del user_contexts[user_id]
            if client:
                await client.cleanup()
            logger.info(f"Cleaned up resources for user {user_id}")